#include "../DWMRICore/Fiber.h"
#include "../DWMRICore/Scalar.h"
#include "../DWMRICore/DTI.h"

#include "StructureTensor.h"

#include <omp.h>

char *fiber_pathname = NULL;
char *dst_pathname = NULL;
char *tempPathname = NULL;

char **confPathname = NULL;
char **vertPathname = NULL;
int w, h, d, scale;
float stepSize, maxLength;


CDTI *pStructureTensor = NULL;
CDTI *pTempTensor = NULL;

float fiber_length = 0.0f;
int kernel_size = 0;
int mode = 0;		/* 0 for dti, 1 for hardi */

void PrintUsage()
{
	printf("Usage: ComputeStructureTensor -fiber <configure_file> -dst <dst_file> -length <fiber_length> -kernel <kernel_size> -mode <mode> [-temp <temp_file>]\n");
}

int ParseArguments(const int argc, char *argv[])
{
	for (int i = 1; i < argc; ++i) {
		if (strcmp(argv[i], "-fiber") == 0) {
			fiber_pathname = argv[++i];
		} else if (strcmp(argv[i], "-dst") == 0) {
			dst_pathname = argv[++i];
		} else if (strcmp(argv[i], "-length") == 0) {
			fiber_length = atof(argv[++i]);
		} else if (strcmp(argv[i], "-kernel") == 0) {
			kernel_size = atoi(argv[++i]);
		} else if (strcmp(argv[i], "-temp") == 0) {
			tempPathname = argv[++i];
		} else if (strcmp(argv[i], "-mode") == 0) {
			i++;
			if (argv[i][0] == 'D' || argv[i][0] == 'd') {
				mode = 0;
			} else if (argv[i][0] == 'H' || argv[i][0] == 'h') {
				mode = 1;
			} else {
				return -1;
			}
			
		} else {
			return -1;
		}
	}

	if (fiber_pathname == NULL || dst_pathname == NULL || fiber_length == 0.0f || kernel_size == 0.0f)
		return -2;

	return 0;
}

void ComputeStructureTensor()
{
    int offset[9];
    offset[ 0] = 0;
    offset[ 1] = -1;
    offset[ 2] = 1;
    offset[ 3] = -w-1;
    offset[ 4] = -w;
    offset[ 5] = -w+1;
    offset[ 6] = w-1;
    offset[ 7] = w;
    offset[ 8] = w+1;
    
    omp_set_num_threads(64);

	float length = fiber_length > maxLength ? maxLength : fiber_length;
	length = length / stepSize;
	int step_delta = int((length * RESAMPLE) / ((PARAM_COUNT-1) / 2));

	float *tensors = NULL;
	if (mode == 1) {
		tensors = new float[w*h*d*30];
		memset(tensors, 0, sizeof(float)*w*h*d*30);

		for (int z = 0; z < d; ++z) {
			for (int y = 0; y < h; ++y) {           
				for (int x = 0; x < w; ++x) {
					int index = (z * h + y) * w  + x;
					tensors[index*30+0] = 1.0f;
					
					tensors[index*30+1] = 0.0f;
					tensors[index*30+2] = 0.0f;
					tensors[index*30+3] = 0.0f;

					tensors[index*30+4] = 1.0f;
					tensors[index*30+5] = 0.0f;
					tensors[index*30+6] = 0.0f;
					tensors[index*30+7] = 1.0f;
					tensors[index*30+8] = 0.0f;
					tensors[index*30+9] = 1.0f;
				}
			}
		}
	}

	for (int z = 1; z < d - 1; ++z) {
        CFiber **fibers[3];
        int *count[3];
        int cw, ch;
        
        ReadFibers(confPathname[z-1], vertPathname[z-1], &(fibers[0]), &(count[0]), cw, ch);
        ReadFibers(confPathname[z+0], vertPathname[z+0], &(fibers[1]), &(count[1]), cw, ch);
		ReadFibers(confPathname[z+1], vertPathname[z+1], &(fibers[2]), &(count[2]), cw, ch);
		
		for (int y = 1; y < h-1; ++y) {
			//printf("y = %d\n", y);
#pragma omp parallel for            
			for (int x = 1; x < w-1; ++x) {
				int index = y * w + x;

                for (int k = 0; k < count[1][index]; ++k) {
                    CFiber *f[27];
                    f[0] = &(fibers[1][index][k]);
                    CVector3F dir = f[0]->m_seed.dir;
                    int id = 1;
					for (int l = 0; l < 3; ++l) {
						for (int i = 0; i < 9; ++i) {
							if (l == 1 && i == 0)
								continue;
							float max_v = -1.0f;
							int max_index;
							int temp_index = index + offset[i];
							int size = count[l][temp_index];
							for (int j = 0; j < size; ++j) {
								float v = InnerProduct(dir, fibers[l][temp_index][j].m_seed.dir);
								if (fabs(v) > max_v) {
									max_v = fabs(v);
									max_index = j;
									if (l == 0) {
										f[9+i] = &(fibers[l][temp_index][j]);
									} else if (l == 1) {
										f[i] = &(fibers[l][temp_index][j]);
									} else {
										f[18+i] = &(fibers[l][temp_index][j]);
									}
								}                            
							}
						}
					}

					float st[7] = {1.0f, 1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f};
					StructureTensor(f, 27, st, step_delta);
					if (mode == 0) {
						pTempTensor->SetDTI(x, y, z, st);
					} else {
						int id = ((z * h + y) * w + x) * 3 + k;
						tensors[id*10] = 1.0f;
						tensors[id*10+1] = dir.m_x;
						tensors[id*10+2] = dir.m_y;
						tensors[id*10+3] = dir.m_z;
						tensors[id*10+4] = st[1];
						tensors[id*10+5] = st[2];
						tensors[id*10+6] = st[3];
						tensors[id*10+7] = st[4];
						tensors[id*10+8] = st[5];
						tensors[id*10+9] = st[6];
					}
                }
			}
		}

        for (int i = 0; i < w * h; ++i) {
            SafeDeleteArray(fibers[0][i]);
            SafeDeleteArray(fibers[1][i]);
            SafeDeleteArray(fibers[2][i]);
        }
	    SafeDeleteArray(fibers[0]);
        SafeDeleteArray(fibers[1]);
        SafeDeleteArray(fibers[2]);	
	    SafeDeleteArray(count[0]);
        SafeDeleteArray(count[1]);
        SafeDeleteArray(count[2]);	

        printf("z = %d\n", z);
        fflush(stdout);
	}

	if (mode == 0 && tempPathname != NULL) {
		pTempTensor->SaveDTIFile(tempPathname);
	} else if (mode == 1 && tempPathname != NULL) {
		FILE *fp = fopen(tempPathname, "wb");
		fwrite(tensors, sizeof(float), w*h*d*30, fp);
		fclose(fp);
	} else {
		;
	}
	
	/* apply a gaussian convolution kernel on the tensor */
	int ks = 2 * kernel_size + 1;
	float *gaussian = new float[ks*ks*ks];
	for (int z = 0; z < ks; ++z) {
		for (int y = 0; y < ks; ++y) {
			for (int x = 0; x < ks; ++x) {
				int index = (z * ks + y) * ks + x;
				float dis = (x - kernel_size) * (x - kernel_size) + (y - kernel_size) * (y - kernel_size) + (z - kernel_size) * (z - kernel_size);
				gaussian[index] = exp(-dis / (0.25f * ks * ks));
			}
		}
	}

	for (int z = 0; z < d; ++z) {
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				int min_x = (x - kernel_size) > 0 ? (x - kernel_size) : 0;
				int max_x = (x + kernel_size) < w ? (x + kernel_size) : (w - 1);
				int min_y = (y - kernel_size) > 0 ? (y - kernel_size) : 0;
				int max_y = (y + kernel_size) < h ? (y + kernel_size) : (h - 1);
				int min_z = (z - kernel_size) > 0 ? (z - kernel_size) : 0;
				int max_z = (z + kernel_size) < d ? (z + kernel_size) : (d - 1);

				float st[7] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};

				if (mode == 0) {
					float temp_st[7];
					float weight = 0.0f;
					for (int cz = min_z; cz <= max_z; ++cz) {
						for (int cy = min_y; cy <= max_y; ++cy) {
							for (int cx = min_x; cx <= max_x; ++cx) {
								pTempTensor->GetDTIFast(cx, cy, cz, temp_st);
								int kx = kernel_size - (x - cx);
								int ky = kernel_size - (y - cy);
								int kz = kernel_size - (z - cz);
								int index = (kz * ks + ky) * ks + kx;

								for (int i = 0; i < 7; ++i) {
									st[i] += gaussian[index] * temp_st[i];
								}
								weight += gaussian[index];
							}
						}
					}

					for (int i = 0; i < 7; ++i) {
						st[i] = st[i] / weight;
					}
				} else {
					int index = (z * h + y) * w + x;
					int count = 0;
					for (int k = 0; k < 3; ++k) {
						if (tensors[(index*3+k)*10] > 0.5f) {
							CVector3F dir;
							dir.m_x = tensors[(index*3+k)*10+1];
							dir.m_y = tensors[(index*3+k)*10+2];
							dir.m_z = tensors[(index*3+k)*10+3];

							float inner_st[7] = {0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 0.0f};
							float temp_st[7];
							float weight = 0.0f;
							for (int cz = min_z; cz <= max_z; ++cz) {
								for (int cy = min_y; cy <= max_y; ++cy) {
									for (int cx = min_x; cx <= max_x; ++cx) {
										
										int closeID = 0;
										float closeV = 0.0f;
										int temp_index = (cz * h + cy) * w + cx;
										for (int m = 0; m < 3; ++m) {
											if (tensors[(temp_index*3+m)*10] > 0.5f) {
												CVector3F temp_dir;
												temp_dir.m_x = tensors[(temp_index*3+m)*10+1];
												temp_dir.m_y = tensors[(temp_index*3+m)*10+2];
												temp_dir.m_z = tensors[(temp_index*3+m)*10+3];
												if (fabs(InnerProduct(dir, temp_dir)) > closeV) {
													closeV = fabs(InnerProduct(dir, temp_dir));
													closeID = m;
												}
											}
										}

										temp_st[0] = 1.0f;
										temp_st[1] = tensors[(temp_index*3+closeID)*10+4];
										temp_st[2] = tensors[(temp_index*3+closeID)*10+5];
										temp_st[3] = tensors[(temp_index*3+closeID)*10+6];
										temp_st[4] = tensors[(temp_index*3+closeID)*10+7];
										temp_st[5] = tensors[(temp_index*3+closeID)*10+8];
										temp_st[6] = tensors[(temp_index*3+closeID)*10+9];

										int kx = kernel_size - (x - cx);
										int ky = kernel_size - (y - cy);
										int kz = kernel_size - (z - cz);
										int index = (kz * ks + ky) * ks + kx;

										for (int i = 0; i < 7; ++i) {
											inner_st[i] += gaussian[index] * temp_st[i];
										}
										weight += gaussian[index];
									}
								}
							}

							for (int i = 0; i < 7; ++i) {
								st[i] += inner_st[i] / weight;
							}

							count++;
						}
					}

					for (int i = 0; i < 7; ++i) {
						st[i] = st[i] / (float)count;
					}
				}

				pStructureTensor->SetDTI(x, y, z, st);
			}
		}
	}

	SafeDeleteArray(gaussian);
	SafeDeleteArray(tensors);
}

int main(int argc, char *argv[])
{
	if (ParseArguments(argc, argv) != 0) {
		PrintUsage();
		return 0;
	}
	
	/* read the configure file */
	ReadConfigureFile(fiber_pathname, &confPathname, &vertPathname, w, h, d, scale, stepSize, maxLength);
	
	/* compute structure tensor for each voxel */
	pTempTensor = new CDTI();
	pTempTensor->CreateDTI(w, h, d);
	for (int z = 0; z < d; ++z) {
		for (int y = 0; y < h; ++y) {
			for (int x = 0; x < w; ++x) {
				float dti[7] = {1.0f, 1.0f, 0.0f, 0.0f, 1.0f, 0.0f, 1.0f};
				pTempTensor->SetDTI(x, y, z, dti);
			}
		}
	}

	pStructureTensor = new CDTI();
	pStructureTensor->CreateDTI(w, h, d);

	ComputeStructureTensor();
	
    /* save the result */
    pStructureTensor->SaveDTIFile(dst_pathname);
	
	for (int z = 0; z < d; ++z) {
		SafeDeleteArray(confPathname[z]);
		SafeDeleteArray(vertPathname[z]);
	}
	SafeDeleteArray(confPathname);
	SafeDeleteArray(vertPathname);

	SafeDelete(pStructureTensor);
	SafeDelete(pTempTensor);
}
